#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SQLite数据库导入工具
用于将CSV数据导入到SQLite数据库中
"""

import os
import csv
import sqlite3
import sys

class SQLiteImporter:
    def __init__(self, db_path='data/xhs_database.db'):
        """初始化SQLite导入工具"""
        # 确保数据目录存在
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        
        self.db_path = db_path
        self.conn = None
        self.cursor = None

    def connect(self):
        """连接到SQLite数据库"""
        try:
            self.conn = sqlite3.connect(self.db_path)
            self.cursor = self.conn.cursor()
            print(f"已连接到数据库: {self.db_path}")
            return True
        except sqlite3.Error as e:
            print(f"连接数据库时出错: {e}")
            return False

    def create_tables(self):
        """创建数据库表结构"""
        try:
            # 创建笔记表
            self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS notes (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                title TEXT,
                author TEXT,
                content TEXT,
                edit_time TEXT,
                likes INTEGER,
                collects INTEGER,
                comments INTEGER,
                collect_time TEXT,
                url TEXT UNIQUE,
                keyword TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
            ''')
            
            # 创建评论表
            self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS comments (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id TEXT,
                comment_text TEXT,
                comment_time TEXT,
                collect_time TEXT,
                article_url TEXT,
                keyword TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (article_url) REFERENCES notes (url)
            )
            ''')
            
            # 创建索引以优化查询
            self.cursor.execute('CREATE INDEX IF NOT EXISTS idx_notes_keyword ON notes (keyword)')
            self.cursor.execute('CREATE INDEX IF NOT EXISTS idx_notes_url ON notes (url)')
            self.cursor.execute('CREATE INDEX IF NOT EXISTS idx_comments_article_url ON comments (article_url)')
            self.cursor.execute('CREATE INDEX IF NOT EXISTS idx_comments_keyword ON comments (keyword)')
            
            self.conn.commit()
            print("数据库表结构创建完成")
            return True
        except sqlite3.Error as e:
            print(f"创建表结构时出错: {e}")
            return False

    def import_notes(self, csv_file='data/notes/xhs_python_notes.csv'):
        """导入笔记数据"""
        if not os.path.exists(csv_file):
            print(f"笔记文件不存在: {csv_file}")
            return False
            
        try:
            with open(csv_file, 'r', encoding='utf-8-sig') as f:
                reader = csv.reader(f)
                headers = next(reader)  # 跳过表头
                
                # 检查表头是否符合期望
                expected_headers = ['title', 'author', 'content', 'edit_time', 'likes', 'collects', 'comments', 'collect_time', 'url', 'keyword']
                if not all(h in headers for h in expected_headers):
                    print(f"CSV文件表头不符合预期: {headers}")
                    return False
                
                # 准备插入语句
                insert_sql = '''
                INSERT OR IGNORE INTO notes 
                (title, author, content, edit_time, likes, collects, comments, collect_time, url, keyword) 
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                '''
                
                # 批量插入数据
                count = 0
                for row in reader:
                    if len(row) >= 10:
                        # 转换数字字段
                        likes = int(row[4]) if row[4].isdigit() else 0
                        collects = int(row[5]) if row[5].isdigit() else 0
                        comments = int(row[6]) if row[6].isdigit() else 0
                        
                        # 插入数据
                        self.cursor.execute(insert_sql, (
                            row[0],          # title
                            row[1],          # author
                            row[2],          # content
                            row[3],          # edit_time
                            likes,           # likes
                            collects,        # collects
                            comments,        # comments
                            row[7],          # collect_time
                            row[8],          # url
                            row[9]           # keyword
                        ))
                        count += 1
                        
                        # 每100条提交一次
                        if count % 100 == 0:
                            self.conn.commit()
                            print(f"已导入 {count} 条笔记数据")
                
                # 最终提交
                self.conn.commit()
                print(f"笔记数据导入完成，共导入 {count} 条记录")
                return True
                
        except Exception as e:
            print(f"导入笔记数据时出错: {e}")
            return False

    def import_comments(self, csv_file='data/user/xhs_comments.csv'):
        """导入评论数据"""
        if not os.path.exists(csv_file):
            print(f"评论文件不存在: {csv_file}")
            return False
            
        try:
            with open(csv_file, 'r', encoding='utf-8-sig') as f:
                reader = csv.reader(f)
                headers = next(reader)  # 跳过表头
                
                # 检查表头是否符合期望
                expected_headers = ['user_id', 'comment_text', 'comment_time', 'collect_time', 'article_url', 'keyword']
                if not all(h in headers for h in expected_headers):
                    print(f"CSV文件表头不符合预期: {headers}")
                    return False
                
                # 准备插入语句
                insert_sql = '''
                INSERT INTO comments 
                (user_id, comment_text, comment_time, collect_time, article_url, keyword) 
                VALUES (?, ?, ?, ?, ?, ?)
                '''
                
                # 批量插入数据
                count = 0
                for row in reader:
                    if len(row) >= 6:
                        # 插入数据
                        self.cursor.execute(insert_sql, (
                            row[0],          # user_id
                            row[1],          # comment_text
                            row[2],          # comment_time
                            row[3],          # collect_time
                            row[4],          # article_url
                            row[5]           # keyword
                        ))
                        count += 1
                        
                        # 每500条提交一次
                        if count % 500 == 0:
                            self.conn.commit()
                            print(f"已导入 {count} 条评论数据")
                
                # 最终提交
                self.conn.commit()
                print(f"评论数据导入完成，共导入 {count} 条记录")
                return True
                
        except Exception as e:
            print(f"导入评论数据时出错: {e}")
            return False

    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
            print("数据库连接已关闭")

def main():
    """主函数"""
    # 创建导入器实例
    importer = SQLiteImporter()
    
    # 连接数据库
    if not importer.connect():
        return
    
    try:
        # 创建表结构
        if not importer.create_tables():
            return
        
        # 导入笔记数据
        notes_path = 'data/notes/xhs_python_notes.csv'
        if len(sys.argv) > 1:
            notes_path = sys.argv[1]
        print(f"导入笔记数据: {notes_path}")
        importer.import_notes(notes_path)
        
        # 导入评论数据
        comments_path = 'data/user/xhs_comments.csv'
        if len(sys.argv) > 2:
            comments_path = sys.argv[2]
        print(f"导入评论数据: {comments_path}")
        importer.import_comments(comments_path)
        
        print("数据导入完成")
        
    finally:
        # 关闭数据库连接
        importer.close()

if __name__ == "__main__":
    main() 